#include "path_generate/p2p/jps.h"
#include "core/utils/graphis_utils.h"
#include <queue>
#include <unordered_map>
#include <vector>
#include <algorithm>  // for std::reverse

struct CompareNodes {
    bool operator()(const GridNodePtr& lhs, const GridNodePtr& rhs) const {
        return lhs->getF() > rhs->getF() 
            || (lhs->getF() == rhs->getF() && lhs->getH() > rhs->getH());
    }
};


JPSPlanner::JPSPlanner(){ 
    debugMapPub_ = nh_.advertise<nav_msgs::OccupancyGrid>("/grid_expend_map", 1, true);
    jumpPointPub_ = nh_.advertise<visualization_msgs::MarkerArray>("/jps/jump_points", 1, true);
}


// 路径规划
void JPSPlanner::genGridPath(int sx, int sy, int gx, int gy, vector<pair<int, int>>& points){
    LOG(INFO) <<  "start: (" << sx << "," << sy
        << "), goal: (" << gx << "," << gy << ")";

    if(debug_){ //调试模式
        expend_points_.clear();
        jump_points_.clear();
        debugGrid_ = getOccupancyGrid();
    }
    int width = getWidth();
    int height = getHeight();

    // 设置起点、终点
    GridNodePtr nodeStartPtr = std::make_shared<GridNode>(sx, sy, 0, 0);
    nodeStartPtr->setIdx(width);


    GridNodePtr nodeGoalPtr = std::make_shared<GridNode>(gx, gy, 0, 0);
    nodeGoalPtr->setIdx(width);

    // 初始化列表
    std::priority_queue<GridNodePtr, std::vector<GridNodePtr>, CompareNodes> openList;
    std::unordered_map<int, float> closedList;  //存储G值
    nodeStartPtr->updateH(*nodeGoalPtr);
    openList.push(nodeStartPtr);
    
    int iteration = 0;
    while(!ros::isShuttingDown() && !openList.empty()){
        // 获取优先级最高节点
        auto nodeCurr = openList.top();
        openList.pop();
        int idxPred = nodeCurr->getIdx();

        iteration++;
        if(debug_){ //调试模式
            cout << "iteration: " << iteration << ", top point ("
                << nodeCurr->getX() << "," << nodeCurr->getY() << ")"
                << ", g: " << nodeCurr->getG()
                << ", h: " << nodeCurr->getH()
                << ", f: " << nodeCurr->getF()
                << endl;
            // if(iteration >= 20){
            //     break;
            // }
        }

        if(closedList.find(idxPred) != closedList.end()){
            continue;
        }
        closedList[idxPred] = nodeCurr->getG();

        if(debug_){ //调试模式
            expend_points_.insert(idxPred);
            // cout << "idxPred: " << idxPred 
            //     << ", x: " << nodeCurr->getX() 
            //     << ", y: " << nodeCurr->getY()
            //     << ", cost: " << nodeCurr->getF()
            //     << endl;
        }

        // 发现目标节点
        if(*nodeCurr == *nodeGoalPtr){
            LOG(INFO) << "find goal. iteration: " << iteration;
            vector<pair<int, int>> origin_points;
            while(nodeCurr->getPred() != nullptr){
                origin_points.emplace_back(make_pair(nodeCurr->getX(), nodeCurr->getY()));
                nodeCurr = nodeCurr->getPred();
            } 
            // 添加起点
            origin_points.emplace_back(make_pair(nodeCurr->getX(), nodeCurr->getY()));
            std::reverse(origin_points.begin(), origin_points.end());   
            LOG(INFO) << "origin_points size: " << origin_points.size();

            // 填充线段之间点集
            for(int i=0; i<origin_points.size()-1; i++){
                int start_x = origin_points[i].first;
                int start_y = origin_points[i].second;
                int goal_x = origin_points[i+1].first;
                int goal_y = origin_points[i+1].second;

                int dx = goal_x - start_x;
                int dy = goal_y - start_y;
                int nums = 0;
                if(dx == 0){
                    nums = abs(dy);
                    dy = dy > 0 ? 1 : -1;
                }
                else if(dy == 0){
                    nums = abs(dx);
                    dx = dx > 0 ? 1 : -1;                    
                }
                else{
                    // 此时abs(dx) == (dy)
                    nums = abs(dx);
                    dx = dx > 0 ? 1 : -1;  
                    dy = dy > 0 ? 1 : -1;
                }
                // cout << "i: " << i 
                //     << ", start: (" << start_x << "," << start_y 
                //     << "), goal: (" << goal_x << "," << goal_y << ")"
                //     << ", nums: " << nums << ", dx: " << dx << ", dy: " << dy
                //     << endl; 

                for(int n=0; n<nums; n++){
                    int new_x = start_x + n * dx;
                    int new_y = start_y + n * dy;
                    // cout << "n: " << n << ", new_x: " << new_x << ", new_y: " << new_y << endl;
                    points.emplace_back(make_pair(new_x, new_y));
                }
            }
            if(origin_points.size() > 1) //添加终点
                points.emplace_back(make_pair(origin_points.back().first, origin_points.back().second));

            break;
        }


        vector<pair<int, int>> prune_dirs = pruneNeighbor(nodeCurr);
        for(int i = 0; i < prune_dirs.size(); i++) {
            auto nodeJump = jump(nodeCurr, prune_dirs[i].first, prune_dirs[i].second, nodeGoalPtr);
            if(nodeJump == nullptr) {
                continue;
            }
            int idxJump = nodeJump->getIdx();
            // nodeJump->setPred(nodeCurr);
            // nodeJump->updateG();
            // nodeJump->updateH(*nodeGoalPtr);

            if(debug_){ //调试模式
                jump_points_.insert(idxJump);
                // cout << "--> iteration: " << iteration << ", add jump point ("
                //     << nodeJump->getX() << "," << nodeJump->getY() << ")"
                //     << ", g: " << nodeJump->getG()
                //     << ", h: " << nodeJump->getH()
                //     << ", f: " << nodeJump->getF()
                //     << endl;
            }

            // 障碍物检测 是否在闭列表
            if(closedList.find(idxJump) != closedList.end()){
                continue;
            }

            // 添加到开列表中
            nodeJump->setPred(nodeCurr);
            nodeJump->updateG();
            nodeJump->updateH(*nodeGoalPtr);
            openList.push(nodeJump);   

        }
    }

    if(debug_){ //调试模式
        LOG(INFO) << "openList size: " << openList.size() 
            << ", closedList size: " << closedList.size();
        for(auto idx : expend_points_){
            if(idx >=0 && idx < debugGrid_.data.size()){
                debugGrid_.data[idx] = 30;
            }
        }
        debugMapPub_.publish(debugGrid_);

        // 发布跳点数据
        vector<geometry_msgs::Point> marker_points;
        for(auto idx : jump_points_){
            int mx,my;
            indexToMap(idx, mx, my);
            geometry_msgs::Point point;
            mapToWorld(mx, my, point.x, point.y);
            marker_points.emplace_back(point);
        }
        visualization_msgs::MarkerArray markerArray_msg;
        visualization_msgs::Marker marker_msg
            = rvizMarker_.genSphereListMarker(marker_points, 0.05);
        rvizMarker_.setMarkerHeader(marker_msg, "map", "jps");
        rvizMarker_.setMarkerId(marker_msg, 0);
        rvizMarker_.setMarkerColor(marker_msg, rvizMarker_.colorBlue(1.0));
        markerArray_msg.markers.push_back(marker_msg);
        jumpPointPub_.publish(markerArray_msg);
    }
}

// 裁剪邻居节点
vector<pair<int, int>> JPSPlanner::pruneNeighbor(GridNodePtr curr_node){
    vector<pair<int, int>> prune_dirs;
    if(!curr_node->getPred()){ 
        // 如果是起点，不需要裁剪
        for(int i = 0; i < GridNode::dir; i++) {
            prune_dirs.emplace_back(make_pair(GridNode::dx[i], GridNode::dy[i]));
        }
    }
    else{
        int pred_dir_x, pred_dir_y;
        curr_node->getDir(pred_dir_x, pred_dir_y);
        int x = curr_node->getX();
        int y = curr_node->getY();

        // 若有父节点，则依据父节点方向进行裁剪
        // 对角方向运动
        if (pred_dir_x != 0 && pred_dir_y != 0) {
            // 自然邻居：保持对角方向、水平和垂直方向
            if(valueIsValid(x+pred_dir_x, y)){
                prune_dirs.emplace_back(make_pair(pred_dir_x, 0));
            }
            if(valueIsValid(x, y+pred_dir_y)){
                prune_dirs.emplace_back(make_pair(0, pred_dir_y));
            }
            if(valueIsValid(x+pred_dir_x, y+pred_dir_y)){
                prune_dirs.emplace_back(make_pair(pred_dir_x, pred_dir_y));
            }
            
            // 强迫邻居：
            if(!valueIsValid(x, y-pred_dir_y)){
                prune_dirs.emplace_back(make_pair(pred_dir_x, -pred_dir_y));
            }
            if(!valueIsValid(x-pred_dir_x, y)){
                prune_dirs.emplace_back(make_pair(-pred_dir_x, pred_dir_y));
            }
        }
        else{
            if (pred_dir_y == 0) {  //水平方向
                // 自然邻居：继续水平方向
                if(valueIsValid(x+pred_dir_x, y)){
                    prune_dirs.emplace_back(make_pair(pred_dir_x, 0));
                }
                // 强迫邻居：检查上、下两个方向
                if(!valueIsValid(x, y-1)){
                    prune_dirs.emplace_back(make_pair(pred_dir_x, -1));
                }   
                if(!valueIsValid(x, y+1)){
                    prune_dirs.emplace_back(make_pair(pred_dir_x, 1));
                }
            }
            else if (pred_dir_x == 0) { //垂直方向
                // 自然邻居：继续垂直方向
                if(valueIsValid(x, y + pred_dir_y)){
                    prune_dirs.emplace_back(make_pair(0, pred_dir_y));
                }
                // 强迫邻居：检查左右两个方向
                if(!valueIsValid(x-1, y)){
                    prune_dirs.emplace_back(make_pair(-1, pred_dir_y));
                }
                if(!valueIsValid(x+1, y)){
                    prune_dirs.emplace_back(make_pair(1, pred_dir_y));
                }
            }            
        }
    }
    return prune_dirs;
}

// 检查是否为强迫节点
bool JPSPlanner::hasForcedNeighbor(GridNodePtr curr_node, int dir_x, int dir_y){
    int x = curr_node->getX();
    int y = curr_node->getY();

    // 水平或垂直方向移动
    if(dir_x == 0 || dir_y == 0) {
        // 水平移动：检查上下方向的强迫邻居
        if (dir_x != 0) {
            // 如果上方被障碍物挡住但右上/左上方向通行，则构成强迫邻居
            if(!valueIsValid(x, y+1) && valueIsValid(x + dir_x, y+1)){
                return true;
            }
            if(!valueIsValid(x, y-1) && valueIsValid(x + dir_x, y-1)){
                return true;
            }
        }
        // 垂直移动：检查左右方向的强迫邻居
        else {
            if(!valueIsValid(x+1, y) && valueIsValid(x+1, y+dir_y)){
                return true;
            }
            if(!valueIsValid(x-1, y) && valueIsValid(x-1, y+dir_y)){
                return true;
            }
        }
    }
    // 对角线方向移动
    else {
        if(!valueIsValid(x-dir_x, y) && valueIsValid(x-dir_x, y+dir_y)){
            return true;
        }
        if(!valueIsValid(x, y-dir_y) && valueIsValid(x+dir_x, y-dir_y)){
            return true;
        }
    }
    return false;
}

// 跳点搜索
GridNodePtr JPSPlanner::jump(GridNodePtr curr_node, int dir_x, int dir_y, GridNodePtr goal){
    int next_x = curr_node->getX() + dir_x;
    int next_y = curr_node->getY() + dir_y;

    if(!valueIsValid(next_x, next_y)){
        return nullptr;
    }

    int width = getWidth();
    auto nodeJump = std::make_shared<GridNode>(next_x, next_y, curr_node->getG(), 0); 
    int idxJump = nodeJump->setIdx(width); 
    nodeJump->setDir(dir_x, dir_y);
    
    if(debug_){ //调试模式
        expend_points_.insert(idxJump);
        // cout << "dir_x: " << dir_x << ", dir_y: " << dir_y << ", add jump point ("
        //     << nodeJump->getX() << "," << nodeJump->getY() << ")"
        //     << endl;
    }

    // 检测是否达到目标点
    if(*nodeJump == *goal){
        return nodeJump;
    }

    // 如果后继节点是强迫节点
    if(hasForcedNeighbor(nodeJump, dir_x, dir_y)){
        return nodeJump;
    }

    // 检测是否为对角线方向
    if (dir_x != 0 && dir_y != 0) {
        // 检测对角线的垂直和水平分量
        if(jump(nodeJump, dir_x, 0, goal) || jump(nodeJump, 0, dir_y, goal)){
            return nodeJump;
        }
    }

    // 沿着当前方向继续递归搜索
    return jump(nodeJump, dir_x, dir_y, goal);
}